import os
import time

import yaml
import subprocess


def get_command_output(command):
    result = subprocess.run(
        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
    )
    # print(result)
    output_lines = result.stdout.strip().split('\n')
    last_line = output_lines[-1] if output_lines else None
    return last_line


DEVICE_NAME = "2080Ti"
EXP_DIR = './record'
os.makedirs(EXP_DIR, exist_ok=True)
# fmt: off
models = ['resnet50','vgg16','resnext26-m']
evaluation_methods = ["ACC", "ASR", "ACAC", "ACTC", "BD", "NU", "NSense", "MPath", "NCoverage"]
# models = ['botnet', "cct", 'ceit', 'coatnet', 'cpvt', 'cvt', 'densenet', 'levit', 'mlpmixer', 'mobilevit', 'poolformer',
# 'PreActResNet50', 'pvtv2', 'resnet50', 'ResNeXt29', 'swin_t', 'vgg16', 'vit', 'ViTAE_basic_small', 'wideresnet', 'shufflenet']
# fmt: on


def modify_corrupt_config(attack, corruption_name):
    yaml_path = r'config/' + attack + '.yaml'
    with open(yaml_path, encoding='utf-8') as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
        f.close()
    config['corruption_name'] = corruption_name
    with open(yaml_path, 'w') as file:
        yaml.dump(config, file)
        file.close()


def modify_system_config(attack, decoder_type, resize_type):
    yaml_path = r'config/' + attack + '.yaml'
    with open(yaml_path, encoding='utf-8') as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    config['decoder_type'] = decoder_type
    config['resize_type'] = resize_type
    with open(yaml_path, 'w') as file:
        yaml.dump(config, file)


def deal_all_attack():
    """
    a table like 
    none ["ACC", "ASR", "ACAC", "ACTC", "BD", "NU", "NSense", "MPath", "NCoverage"]
    fgsm
    bim
    pgd
    ...
    """
    ### adversarial attack
    attack_methods = ['fgsm', 'bim', 'pgd']
    for model in models:
        lines = []
        headline = "none," + ",".join(evaluation_methods)
        lines.append(headline)
        for attack in attack_methods:
            attack_line = attack.upper()
            for evaluation_method in evaluation_methods:
                instruction = (
                    "python main.py --evaluation_method="
                    + evaluation_method
                    + " --model="
                    + model
                    + " --attack="
                    + attack
                )
                print("\n---------------------------------------\n")
                print(instruction)
                time1 = time.time()
                # os.system(instruction)
                output = get_command_output(instruction)
                print(output)
                value = output[len(evaluation_method) + 1 :].strip()
                print(f"value is {value}")
                time2 = time.time()
                # fmt: off
                print(
                    "evaluation_method: " + evaluation_method + ", model: " + model + ", attack: "+attack,
                    ", time: " + str(time2 - time1),
                )
                # fmt: on
                attack_line += ',' + value

            lines.append(attack_line)
        table_name = model + '_adversarial_' + DEVICE_NAME + '.csv'
        f_table = open(os.path.join(EXP_DIR, table_name), 'w')
        f_table.write("\n".join(lines))


def deal_all_corrupt():
    """
    a table like 
    none ["ACC", "ASR", "ACAC", "ACTC", "BD", "NU", "NSense", "MPath", "NCoverage"]
    gaussian_noise
    shot_noise
    impulse_noise
    ...
    """
    ### adversarial attack
    category_attack = 'corAttack'
    corruption_name = "gaussian_noise, shot_noise, impulse_noise,defocus_blur, glass_blur,zoom_blur,frost, fog, brightness, contrast,elastic_transform, pixelate, jpeg_compression,speckle_noise, gaussian_blur, spatter, saturate"
    # corruption_name = "spatter"
    attack_methods = corruption_name.replace(" ", "").split(',')
    print(f"corrupt is {attack_methods}")
    for model in models:
        lines = []
        headline = "none," + ",".join(evaluation_methods)
        lines.append(headline)
        table_name = model + '_corAttack_' + DEVICE_NAME + '.csv'
        f_table = open(os.path.join(EXP_DIR, table_name), 'w')
        f_table.write(headline + '\n')
        f_table.flush()
        for attack in attack_methods:
            modify_corrupt_config(category_attack, corruption_name=attack)
            attack_line = attack
            for evaluation_method in evaluation_methods:
                instruction = (
                    "python main.py --evaluation_method="
                    + evaluation_method
                    + " --model="
                    + model
                    + " --attack="
                    + category_attack
                )
                print("\n---------------------------------------\n")
                print(instruction)
                time1 = time.time()
                # os.system(instruction)
                output = get_command_output(instruction)
                print(output)
                value = output[len(evaluation_method) + 1 :].strip()
                print(f"value is {value}")
                time2 = time.time()
                # fmt: off
                print(
                    "evaluation_method: " + evaluation_method + ", model: " + model + ", attack: "+attack,
                    ", time: " + str(time2 - time1),
                )
                # fmt: on
                attack_line += ',' + value

            lines.append(attack_line)
            f_table.write(attack_line + '\n')
            f_table.flush()
        # f_table.write("\n".join(lines))
        f_table.close()


def deal_all_system():
    """
    a table like 
    none ["ACC", "ASR", "ACAC", "ACTC", "BD", "NU", "NSense", "MPath", "NCoverage"]
    opencv-nearest
    shot_noise
    impulse_noise
    ...
    """
    ### adversarial attack
    category_attack = 'sysAttack'

    decoder_type = ['pil', 'opencv', 'ffmpeg']
    resize_type = [
        'pil-bilinear',
        'pil-nearest',
        'pil-box',
        'pil-hamming',
        'pil-cubic',
        'pil-lanczos',
        'opencv-nearest',
        'opencv-bilinear',
        'opencv-area',
        'opencv-cubic',
        'opencv-lanczos',
    ]

    attack_methods = resize_type + decoder_type
    decoder_types = ['pil'] * 6 + ['opencv'] * 5 + decoder_type
    resize_types = [
        'pil-bilinear',
        'pil-nearest',
        'pil-box',
        'pil-hamming',
        'pil-cubic',
        'pil-lanczos',
        'opencv-nearest',
        'opencv-bilinear',
        'opencv-area',
        'opencv-cubic',
        'opencv-lanczos',
        'pil-bilinear',
        'pil-bilinear',
        'pil-bilinear',
    ]
    print(f"systemAttack is {attack_methods}")
    for model in models:
        lines = []
        headline = "none," + ",".join(evaluation_methods)
        lines.append(headline)
        for i, attack in enumerate(attack_methods):
            if attack != "ffmpeg":
                continue
            modify_system_config(
                category_attack,
                decoder_type=decoder_types[i],
                resize_type=resize_types[i],
            )
            attack_line = attack
            for evaluation_method in evaluation_methods:
                instruction = (
                    "python main.py --evaluation_method="
                    + evaluation_method
                    + " --model="
                    + model
                    + " --attack="
                    + category_attack
                )
                print("\n---------------------------------------\n")
                print(instruction)
                time1 = time.time()
                # os.system(instruction)
                output = get_command_output(instruction)
                print(output)
                value = output[len(evaluation_method) + 1 :].strip()
                print(f"value is {value}")
                time2 = time.time()
                # fmt: off
                print(
                    "evaluation_method: " + evaluation_method + ", model: " + model + ", attack: "+attack,
                    ", time: " + str(time2 - time1),
                )
                # fmt: on
                attack_line += ',' + value

            lines.append(attack_line)
        table_name = model + '_sysAttack_' + DEVICE_NAME + '.csv'
        f_table = open(os.path.join(EXP_DIR, table_name), 'w')
        f_table.write("\n".join(lines))


if __name__ == '__main__':
    # deal_all_attack()
    # deal_all_corrupt()
    deal_all_system()
